import torch
from sklearn.manifold import TSNE
import shelve
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
import numpy as np
from scipy.stats import gaussian_kde
from matplotlib import ticker
from matplotlib import colors
from matplotlib import cm

def remove_outliers(data):
    xdata=[]
    ydata=[]
    for i in range(len(data)):
        xdata.append(data[i][0])
        ydata.append(data[i][1])
    xdata = np.array(xdata)
    ydata = np.array(ydata)


    weight = 1
    aa=5

    q1,q2 = np.percentile(xdata,[aa,100-aa])
    q3,q4 = np.percentile(ydata,[aa,100-aa])
    x_iqr=q2-q1
    y_iqr=q4-q3
    x_low_bound = q1-(x_iqr*weight)
    x_high_bound = q2+(x_iqr*weight)
    y_low_bound = q3-(y_iqr*weight)
    y_high_bound = q4+(y_iqr*weight)

    filtered_index = [True]*len(data)
    xcnt=0
    ycnt=0
    for i in range(len(data)):
        if (data[i][0] >=x_low_bound and data[i][0]<=x_high_bound) == False:
            filtered_index[i]=False
            xcnt+=1
        elif (data[i][1] >=y_low_bound and data[i][1]<=y_high_bound) == False:
            filtered_index[i]=False
            ycnt+=1
            
    print("移除{}个离群点,{}个y离群点".format(xcnt,ycnt))
    return data[filtered_index]

def normalize_2d_array(array, min_val=0, max_val=1):
    '''
        数组归一化
    '''
    xlist=[]
    ylist=[]
    for i in range(len(array)):
        xlist.append(array[i][0])
        ylist.append(array[i][1])
    x_min= min(xlist)
    x_max= max(xlist)
    y_min= min(ylist)
    y_max= max(ylist)
    normalize_x = (xlist-x_min)/(x_max-x_min)
    normalize_y = (ylist-y_min)/(y_max-y_min)
    res=[]
    for i in range(len(array)):
        res.append([normalize_x[i],normalize_y[i]])
    ans = np.array(res)
    return ans


def draw_picture_with_tsne(node_embeddings,save_name):
    """
        输入一个array的二维嵌入向量，
        通过t-sne进行降维
        将结构保存为png
    """

    # 创建包含左右两个子图的布局
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))

    # 使用 t-SNE 进行降维
    tsne = TSNE(n_components=2)
    low_dimension_data = tsne.fit_transform(node_embeddings[0])
    
    # 在左侧子图中绘制第一个图形
    axes[0].scatter(low_dimension_data[:, 0], low_dimension_data[:, 1])
    axes[0].set_title('sample x')

    # 使用 t-SNE 进行降维
    tsne = TSNE(n_components=2)
    low_dimension_data = tsne.fit_transform(node_embeddings[1])

    # 在右侧子图中绘制第二个图形
    axes[1].scatter(low_dimension_data[:, 0], low_dimension_data[:, 1])
    axes[1].set_title('sample y')

    # 调整子图之间的间距
    plt.tight_layout()

    # 显示图形
    plt.show()

    plt.savefig(r'/home/cyw/projects/function_sim_project/basic_script/find_sample_pairs/save_img/tsne/{}.png'.format(save_name))
    # 不关闭的话，画布会一直叠加
    plt.close()

def draw_picture_with_pca(node_embeddings,save_name):
    """
        输入一个array的二维嵌入向量，
        通过t-sne进行降维
        将结构保存为png
    """
    # 设置画图的参数
    dotSize = 35
    myAlpha = 1
    color = "black"
    # 创建包含左右两个子图的布局
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))

    # 使用 PCA 进行降维
    pca = PCA(n_components=2)
    low_dimension_data = pca.fit_transform(node_embeddings[0])
    # remove_outliers_data = remove_outliers(low_dimension_data)
    # normailze_data = normalize_2d_array(remove_outliers_data)
    normailze_data = normalize_2d_array(low_dimension_data)
    # normailze_data = remove_outliers(remove_outliers_data)
    x=normailze_data[:, 0]
    y=normailze_data[:, 1]

    # https://blog.csdn.net/fyfugoyfa/article/details/136198893
    # 核密度估计
    x_and_y = np.vstack([x, y])
    kde = gaussian_kde(x_and_y)
    z = kde(x_and_y)
    idx = z.argsort()
    x, y, z = x[idx], y[idx], z[idx]
    is_cbar = False
    my_cmap = "cool"

    # 在左侧子图中绘制第一个图形
    axes[0].scatter(x, y,s=dotSize,c=z,cmap=my_cmap,alpha = myAlpha)
    # axes[0].scatter(x, y,s=dotSize,alpha = myAlpha)
    axes[0].set_title('sample x')
    axes[0].set_xlim(-0.05,1.05)
    axes[0].set_ylim(-0.05,1.05)
    # if is_cbar:
    #     norm = colors.Normalize(vmin=np.min(z), vmax=np.max(z))
    #     cbar = plt.colorbar(cm.ScalarMappable(norm=norm, cmap=my_cmap), ax=axes[0])

    #     cbar.ax.set_ylabel("Density", labelpad=12)
    #     cbar.ax.tick_params(labelsize=12)

    #     labels = cbar.ax.get_xticklabels() + cbar.ax.get_yticklabels()
    #     # [label.set_fontproperties(font_latex2) for label in labels]
    #     [label.set_color('black') for label in labels]
        
    #     tick_locator = ticker.MaxNLocator(nbins=8)
    #     cbar.locator = tick_locator
    #     cbar.update_ticks()

    # 显示网格  虚线和透明度
    axes[0].grid(alpha=0.360, ls="--", which="major", color="#A9A9A9")

    # 使用 PCA 进行降维
    pca = PCA(n_components=2)
    low_dimension_data = pca.fit_transform(node_embeddings[1])
    # remove_outliers_data = remove_outliers(low_dimension_data)
    # normailze_data = normalize_2d_array(remove_outliers_data)
    normailze_data = normalize_2d_array(low_dimension_data)
    # normailze_data = remove_outliers(remove_outliers_data)
    x=normailze_data[:, 0]
    y=normailze_data[:, 1]
    # https://blog.csdn.net/fyfugoyfa/article/details/136198893
    # 核密度估计
    x_and_y = np.vstack([x, y])
    kde = gaussian_kde(x_and_y)
    z = kde(x_and_y)
    idx = z.argsort()
    x, y, z = x[idx], y[idx], z[idx]
    is_cbar = False
    my_cmap = "cool"

    # 在右侧子图中绘制第二个图形
    axes[1].scatter(x, y,s=dotSize,c=z,cmap=my_cmap,alpha = myAlpha)  
    # axes[1].scatter(x, y,s=dotSize,alpha = myAlpha)  
    axes[1].set_title('sample y')
    axes[1].set_xlim(-0.05,1.05)
    axes[1].set_ylim(-0.05,1.05)

    axes[1].grid(alpha=0.360, ls="--", which="major", color="#A9A9A9")


    # 调整子图之间的间距
    plt.tight_layout()

    # 显示图形
    plt.show()

    plt.savefig(r'/home/cyw/projects/function_sim_project/basic_script/find_sample_pairs/save_img/pca/{}.png'.format(save_name))
    # 不关闭的话，画布会一直叠加
    plt.close()

def get_data(name):
    """
        获得样本的函数嵌入
    """

    with shelve.open(basePath) as file:
        res=file["embedding"][0]
    print("{}\n\t样本长度:{}\t样本维度:{}".format(basePath,len(res),len(res[0])))
    return res
    

if __name__ == "__main__":
    names=["functionSim","functionSim_hete","functionSim_zero"]
    prefixs=["x_layer_","y_layer_"]
    for i in [0,1,2,3]:
        for name in names:
            suffix=""
            if name=="functionSim":
                suffix="_cross_True_hete_True"
            elif name=="functionSim_hete":
                suffix="_cross_False_hete_True"
            elif name=="functionSim_zero":
                suffix="_cross_False_hete_False"

            data=[]
            for preifx in prefixs:
                tarName="{}{}{}".format(preifx,i,suffix)
                basePath=r"/home/cyw/projects/function_sim_project/all_data/functionEmbedding/{}".format(tarName)
                node_embeddings=get_data(basePath)
                data.append(node_embeddings)
            draw_picture_with_pca(data,"{}{}".format(i,suffix))
            # draw_picture_with_tsne(data,"{}{}".format(i,suffix))
            